#include <stdio.h>
#include <string.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT *tcti_context = NULL;
TSS2_SYS_CONTEXT *sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;

static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

// Check if the return code indicates success, clean up if not
#define CHECK_PASSED(rc)                                                       \
    {                                                                          \
        if ((rc) != TPM2_RC_SUCCESS) {                                         \
            printf("FAILED! [%s:%u] %s\n", __FUNCTION__, __LINE__,             \
                   Tss2_RC_Decode(rc));                                        \
            cleanup(rc);                                                       \
        } else {                                                               \
            printf("PASSED! [%s:%u]\n", __FUNCTION__, __LINE__);               \
        }                                                                      \
    }

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context,
                         TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size,
                                     tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc =
        Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                               TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush =
                Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n",
                       handle_recorder.handles[i]);
            }
        }
    }
    handle_recorder.count = 0;
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Change hierarchy authorization
TSS2_RC hierarchy_change_auth(TSS2_SYS_CONTEXT *sys_context,
                              TPMI_RH_HIERARCHY hierarchy, const char *cur_pwd,
                              const char *new_pwd) {
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(cur_pwd)},
        }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, cur_pwd, strlen(cur_pwd));

    TPM2B_AUTH new_auth = {.size = strlen(new_pwd)};
    memcpy(new_auth.buffer, new_pwd, new_auth.size);

    TSS2_RC rc = Tss2_Sys_HierarchyChangeAuth(
        sys_context, hierarchy, &cmd_auths_array, &new_auth, NULL);
    return rc;
}

TSS2_RC create_primary_storage_key(TSS2_SYS_CONTEXT *sys_context,
                           TPMI_RH_HIERARCHY hierarchy, const char *hierarchy_password,
                           const char *key_password, TPM2_HANDLE *key_handle) {
    // CreatePrimary Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(hierarchy_password)},
        }
    };
    memcpy(cmd_auths.auths[0].hmac.buffer, hierarchy_password, cmd_auths.auths[0].hmac.size);

    TPM2B_SENSITIVE_CREATE in_sensitive = { .sensitive.userAuth = {.size = strlen(key_password)} };
    memcpy(in_sensitive.sensitive.userAuth.buffer, key_password, in_sensitive.sensitive.userAuth.size);

    TPM2B_PUBLIC in_public = {
        // .size = 0,
        .publicArea = {
            .type = TPM2_ALG_RSA,
            .nameAlg = TPM2_ALG_SHA256,
            .objectAttributes = TPMA_OBJECT_DECRYPT 
                              | TPMA_OBJECT_FIXEDTPM 
                              | TPMA_OBJECT_FIXEDPARENT 
                              | TPMA_OBJECT_USERWITHAUTH 
                              | TPMA_OBJECT_RESTRICTED
                              | TPMA_OBJECT_SENSITIVEDATAORIGIN,
            .parameters.rsaDetail = {
                .symmetric = {
                    .algorithm = TPM2_ALG_AES,
                    .keyBits.aes = 128,
                    .mode.aes = TPM2_ALG_CFB,
                },
                .scheme = { .scheme = TPM2_ALG_NULL },
                .keyBits = 2048,
                .exponent = 0,
            },
        },
    };
    TPM2B_DATA outside_info = {.size = 0};
    TPML_PCR_SELECTION creation_pcr = {0};
    
    // CreatePrimary Output Parameters
    TPM2B_PUBLIC out_public = {0};
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths = {0};

    TSS2_RC rc = Tss2_Sys_CreatePrimary(
        sys_context, hierarchy, &cmd_auths, &in_sensitive, &in_public,
        &outside_info, &creation_pcr, key_handle, &out_public,
        &creation_data, &creation_hash, &creation_ticket, &name, &rsp_auths);
    return rc;
}

TSS2_RC create_and_load_aes_key(TSS2_SYS_CONTEXT *sys_context,
                                 TPM2_HANDLE parent_handle, const char *parent_password,
                                 TPM2_HANDLE *key_handle, const char *key_password) {
    TSS2_RC rc;

    // Create Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_create = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(parent_password)},
        }
    };
    memcpy(cmd_auths_create.auths[0].hmac.buffer, parent_password, cmd_auths_create.auths[0].hmac.size);
    
    TPM2B_SENSITIVE_CREATE in_sensitive = { .sensitive.userAuth = {.size = strlen(key_password)} };
    memcpy(in_sensitive.sensitive.userAuth.buffer, key_password, in_sensitive.sensitive.userAuth.size);

    TPM2B_PUBLIC in_public = {
        // .size = 0,
        .publicArea = {
            .type = TPM2_ALG_SYMCIPHER,
            .nameAlg = TPM2_ALG_SHA256,
            .objectAttributes = TPMA_OBJECT_DECRYPT | TPMA_OBJECT_SIGN_ENCRYPT |
                                TPMA_OBJECT_USERWITHAUTH | TPMA_OBJECT_SENSITIVEDATAORIGIN,
            .parameters.symDetail.sym = {
                    .algorithm = TPM2_ALG_AES,
                    .keyBits.aes = 128,
                    .mode.sym = TPM2_ALG_CFB,
                },
            },
    };
    TPM2B_DATA outside_info = {.size = 0};
    TPML_PCR_SELECTION creation_pcr = {0};
    
    // Create Output Parameters
    TPM2B_PRIVATE out_private = {0};
    TPM2B_PUBLIC out_public = {0};
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_create = {0};

    // Create the key
    rc = Tss2_Sys_Create(
        sys_context, parent_handle, &cmd_auths_create, &in_sensitive, &in_public,
        &outside_info, &creation_pcr, &out_private, &out_public,
        &creation_data, &creation_hash, &creation_ticket, &rsp_auths_create);

    // Load Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_load = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(parent_password)},
        }
    };
    memcpy(cmd_auths_load.auths[0].hmac.buffer, parent_password, cmd_auths_load.auths[0].hmac.size);

    // Load Output Parameters
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_load = {0};
    rc = Tss2_Sys_Load(sys_context, parent_handle, &cmd_auths_load,
                       &out_private, &out_public, key_handle, &name, &rsp_auths_load);
    return rc;
}

// Encrypt data using the symmetric key
TSS2_RC aes_encrypt(TSS2_SYS_CONTEXT* sys_context, const char* password, TPM2_HANDLE key_handle,
                TPM2B_MAX_BUFFER* in_data, TPM2B_MAX_BUFFER* out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1,
                                              .auths[0] = {
                                                  .sessionHandle = TPM2_RH_PW,
                                                  .hmac = {.size = strlen(password)},
                                              }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));

    TPMI_ALG_CIPHER_MODE mode = TPM2_ALG_NULL;
    TPM2B_IV in_iv = {.size = 16};

    // EncryptDecrypt Output Parameters
    TPM2B_IV out_iv = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_EncryptDecrypt2(sys_context, key_handle, &cmd_auths_array, in_data, TPM2_NO, mode,
                                 &in_iv, out_data, &out_iv, &rsp_auths_array);
    return rc;
}

// Decrypt data using the symmetric key
TSS2_RC aes_decrypt(TSS2_SYS_CONTEXT* sys_context, const char* password, TPM2_HANDLE key_handle,
                TPM2B_MAX_BUFFER* in_data, TPM2B_MAX_BUFFER* out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1,
                                              .auths[0] = {
                                                  .sessionHandle = TPM2_RH_PW,
                                                  .hmac = {.size = strlen(password)},
                                              }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));

    TPMI_ALG_CIPHER_MODE mode = TPM2_ALG_NULL;
    TPM2B_IV in_iv = {.size = 16};

    // EncryptDecrypt Output Parameters
    TPM2B_IV out_iv = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_EncryptDecrypt2(sys_context, key_handle, &cmd_auths_array, in_data, TPM2_YES, mode,
                                 &in_iv, out_data, &out_iv, &rsp_auths_array);
    return rc;
}

int main(int argc, char *argv[]) {

    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_context);
    CHECK_PASSED(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_PASSED(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_PASSED(rc);

    // Change owner hierarchy password
    TPMI_RH_HIERARCHY hierarchy = TPM2_RH_OWNER;
    const char *cur_pwd = "";
    const char *new_pwd = "";
    rc = hierarchy_change_auth(sys_context, hierarchy, cur_pwd, new_pwd);
    CHECK_PASSED(rc);

    // Create a storage key
    TPM2_HANDLE key_handle;
    const char *hpwd = new_pwd;
    const char *kpwd = "XYZ123";
    rc = create_primary_storage_key(sys_context, hierarchy, hpwd, kpwd, &key_handle);
    CHECK_PASSED(rc);
    handle_recorder.handles[handle_recorder.count++] = key_handle; // Record the key handle
    printf("Created Primary Storage Key\n");

    // Create and load AES key using the storage key
    TPM2_HANDLE aes_key_handle;
    const char *aes_key_pwd = "AesKeyPwd";
    rc = create_and_load_aes_key(sys_context, key_handle, kpwd, &aes_key_handle, aes_key_pwd);
    CHECK_PASSED(rc);
    handle_recorder.handles[handle_recorder.count++] = aes_key_handle; // Record the key handle
    printf("Created and Loaded AES Key using the Storage Key\n");

    // Encrypt data
    const char* plaintext = "This is a test message.";
    printf("Plain Text: %s\n", plaintext);
    TPM2B_MAX_BUFFER encrypt_in_data = {
        .size = strlen(plaintext),
        .buffer = {0},
    };
    memcpy(encrypt_in_data.buffer, plaintext, encrypt_in_data.size);
    TPM2B_MAX_BUFFER encrypt_out_data = {0};
    rc = aes_encrypt(sys_context, aes_key_pwd, aes_key_handle, &encrypt_in_data, &encrypt_out_data);
    CHECK_PASSED(rc);
    printf("Encrypt Data Complete\n");

    // Decrypt data
    TPM2B_MAX_BUFFER decrypt_out_data = {0};
    rc = aes_decrypt(sys_context, aes_key_pwd, aes_key_handle, &encrypt_out_data, &decrypt_out_data);
    CHECK_PASSED(rc);
    printf("Decrypt Data Complete\n");
    printf("Decrypted Text: %.*s\n", decrypt_out_data.size, decrypt_out_data.buffer); // Print decrypted datas

    // Cleanup resources
    cleanup(rc);

    return 0;
}
